import os
import cv2
import numpy as np

# === 路径配置 ===
ANNOTATION_DIR = "/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/abdominal_US/AUS/annotations/train"
OUTPUT_ROOT = "/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/abdominal_US_organs/labelsTr"
os.makedirs(OUTPUT_ROOT, exist_ok=True)

# === 颜色与器官映射 (BGR格式，OpenCV默认) ===
COLOR_TO_ORGAN = {
    (100, 0, 100): "liver",        # 紫色  肝脏
    (0, 255, 255): "kidney",       # 黄色  肾脏 (BGR)
    (255, 0, 0): "pancreas",       # 蓝色  胰腺 (BGR)
    (0, 0, 255): "vessels",        # 红色  血管 (BGR)
    (255, 255, 0): "adrenals",     # 浅蓝  肾上腺 (BGR)
    (0, 255, 0): "gallbladder",    # 绿色  胆囊
    (255, 255, 255): "bones",      # 白色  骨骼
    (255, 0, 255): "spleen",       # 粉色  脾脏
}

# === 遍历标注文件 ===
total_files = 0
organs_found = {name: 0 for name in COLOR_TO_ORGAN.values()}

for file in os.listdir(ANNOTATION_DIR):
    if not file.endswith(".png"):
        continue

    total_files += 1
    path = os.path.join(ANNOTATION_DIR, file)
    img = cv2.imread(path)

    if img is None:
        print(f"⚠️ 无法读取: {path}")
        continue

    # 遍历每种器官颜色
    for color, organ in COLOR_TO_ORGAN.items():
        organ_dir = os.path.join(OUTPUT_ROOT, organ)
        os.makedirs(organ_dir, exist_ok=True)

        mask = cv2.inRange(img, color, color)
        if np.sum(mask) == 0:
            continue  # 没有该器官就跳过

        organs_found[organ] += 1

        out_path = os.path.join(organ_dir, file)
        cv2.imwrite(out_path, mask)

# === 输出统计 ===
print("\n✅ 提取完成！")
print(f"总图片数: {total_files}")
for organ, count in organs_found.items():
    print(f"{organ:12s}: {count} 张包含此器官")
